package cn.edu.sjtu;

import hu.uw.pallergabor.dedexer.Dedexer;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileReader;
import java.io.InputStream;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;

import org.w3c.dom.Document;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;

import cn.edu.sjtu.permission.APIPermissionChecker;
import cn.edu.sjtu.permission.CPPermissionChecker;
import cn.edu.sjtu.permission.IntentPermissionChecker;
import cn.edu.sjtu.permission.PermissionChecker;

public class DedexerPlus {

	public static HashMap<String, HashMap<Long, List<HashMap<String, String>>>> registerMapsForClasses = new HashMap<String, HashMap<Long, List<HashMap<String, String>>>>();
    public static HashMap<String, String> androidMap = new HashMap<String, String>();	// for android class's fields
	
	private HashSet<String> permission_app_required = null;
	private HashSet<String> permission_app_acquired = null;
	
	private PermissionChecker apipc = null;
	private PermissionChecker intpc = null;
	private PermissionChecker cppc = null;

	static {
		// TODO
		androidMap.put("android.provider.CallLog$Calls.CONTENT_URI", "\"content://call_log/calls\"");
		androidMap.put("android.provider.Contacts$Groups.CONTENT_URI", "\"content://contacts/groups\"");
		androidMap.put("android.provider.Contacts$Phones.CONTENT_URI", "\"content://contacts/phones\"");
		androidMap.put("android.provider.Contacts$People.CONTENT_URI", "\"content://contacts/people\"");
		androidMap.put("android.provider.Contacts$GroupMembership.CONTENT_URI", "\"content://contacts/groupmembership\"");
	}

	public DedexerPlus(String cpXlsFilePath, String apiXlsFilePath, String intXlsFilePath) {	
		apipc = new APIPermissionChecker(apiXlsFilePath);
		intpc = new IntentPermissionChecker(intXlsFilePath);
		cppc = new CPPermissionChecker(cpXlsFilePath);
	}

	public void checkAppPermission(String path) {
		registerMapsForClasses = new HashMap<String, HashMap<Long, List<HashMap<String, String>>>>();
		
		System.out.println("decompressing apk file...");
		String root = decompressApp(path);
		
		System.out.println("read permission from XML...");
		readAcquiredPermission(root + "/AndroidManifest.xml");
		
		System.out.println("analysing classes.dex...\n");
		Dedexer dedexer = new Dedexer();
		String[] args = { "-r", "-d", root + "/dedexer", root + "/classes.dex" };
		dedexer.run(args);
		
		scanAllFile(root + "/dedexer");
		
		HashSet<String> more = (HashSet<String>) permission_app_acquired.clone();
		HashSet<String> need = (HashSet<String>) permission_app_required.clone();
		comparePermission(more, need);
		showCheckResult(more, need);
	}

	private void scanAllFile(String path) {
		permission_app_required = new HashSet<String>();
		
		LinkedList list = new LinkedList();
        list.add(new File(path));
        File tmp, file[];
        while (!list.isEmpty()) {
            tmp = (File) list.removeFirst();
            if (tmp.isDirectory()) {
                file = tmp.listFiles();
                if (file == null)
                    continue;
                for (int i = 0; i < file.length; i++) {
                    if (file[i].isDirectory())
                        list.add(file[i]);
                    else {
                        analyseDDXFile(file[i].getAbsolutePath());
                    }
                }
            } else {
                analyseDDXFile(tmp.getAbsolutePath());
            }
        }
        
        apipc.checkPermission(this);
		intpc.checkPermission(this);
		cppc.checkPermission(this);
	}
	
	private String decompressApp(String path) {
		String[] ts = path.split("/");
		StringBuffer sb = new StringBuffer();
		int i;
		for (i = 0; i < ts.length - 1; i++) {
			sb.append(ts[i] + "/");
		}
		sb.append(ts[i].substring(0, ts[i].lastIndexOf(".")) + "/");
		
		File f = new File(sb.toString());
		Process process;  
        try {
			process = Runtime.getRuntime().exec(new String[]{"java", "-jar", "apktool.jar", "d", path, sb.toString()},null, null);
			
			if (process.waitFor() == 0) {
				process = Runtime.getRuntime().exec(new String[]{"jar", "xf", path, "classes.dex"},null, new File(sb.toString()));
			}
			
			if (process.waitFor() == 0) {} 
			
		} catch (Exception e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
        
        return sb.toString();
	}

	private void showCheckResult(HashSet<String> more, HashSet<String> need) {
		System.out.println("\nOverprivilege: ");
		if (more.size() == 0) {
			System.out.println("None");
		} else {
			for (String s : more) {
				System.out.println(s);
			}
		}
		
		System.out.println("\nRequired Permissions: ");
		if (need.size() == 0) {
			System.out.println("None");
		} else {
			for (String s : need) {
				System.out.println(s);
			}
		}
	}
	
	private void comparePermission(HashSet<String> more, HashSet<String> need) {
		HashSet<String> common = new HashSet<String>();
		for (String p : more) {
			if (need.contains(p)) {
				common.add(p);
			}
		}
		for (String p : common) {
			more.remove(p);
			need.remove(p);
		}
	}

	private void readAcquiredPermission(String path) {
		permission_app_acquired = new HashSet<String>();
		DocumentBuilderFactory domfac = DocumentBuilderFactory.newInstance();
		try {
			DocumentBuilder dombuilder = domfac.newDocumentBuilder();
			InputStream is = new FileInputStream(path);
			Document doc = dombuilder.parse(is);
			NodeList nodeList = doc.getElementsByTagName("uses-permission");
			for (int i = 0; i < nodeList.getLength(); i++) {
				Node n = nodeList.item(i);
				NamedNodeMap attributes = n.getAttributes();
				String ts = attributes.getNamedItem("android:name").toString();
				int start = ts.indexOf("\"") + 1, end = ts.indexOf("\"", start);
				PermissionChecker.setPermissionSet(permission_app_acquired, ts.substring(start, end));
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}
	
	public void analyseDDXFile(String path) {
		try {
			BufferedReader br = new BufferedReader(new FileReader(path));
			String s = null;
			long lineNumber = -1;
			String methodName = null;
			String className = null;
			while ((s = br.readLine()) != null) {
				if (s.startsWith(".class")) {
					String[] ccn = s.split(" ");
		            className = ccn[ccn.length - 1].replace('/', '.');
		            
				} else if (s.startsWith(".line")) {
					s = s.trim().substring(6);
					lineNumber = Integer.parseInt(s);
				} else if (s.startsWith(".method")) {
					s = s.trim().substring(8);
					String[] ta = s.split(" ");
                    methodName = className + "." + ta[ta.length - 1].split("[(]")[0];
                    
				} else if (s.contains("invoke-")){
					String[] tss = s.split(";")[0].trim().split("\t");
					String opcode = tss[0];
					tss = tss[1].split("}");
					
					String func = tss[1].substring(1).replace('/', '.');
					// check Api
					if (func.contains("java.lang.reflect.Method.invoke")) { //reflection
						String i = tss[0].substring(1).split(",")[0].substring(1);
						int reg = Integer.parseInt(i);
						apipc.setAppPermission(methodName + "@" + lineNumber + "@" + reg + "@" + func + "@" + path);
					} else {
						apipc.setAppPermission(func);
						apipc.analyse(permission_app_required, func, null);
					}

					// check Content
					// TODO: check super class
					if ((func.endsWith(".query") || func.endsWith(".insert") || func.endsWith(".update") || func.endsWith(".delete")) &&
						 !func.startsWith("java")) {
						String i;
						int reg;
						if (opcode.contains("/range")) {  // {v0..vn}
							i = tss[0].substring(1).split("[.]")[0].substring(1);
							reg = Integer.parseInt(i);
							if (!opcode.equals("invoke-static")) {
								reg++;
							}
						} else { // {va,vb,vc}
							if (!opcode.equals("invoke-static")) {
								i = tss[0].substring(1).split(",")[1].substring(1);
							} else {
								i = tss[0].substring(1).split(",")[0].substring(1);
							}
							reg = Integer.parseInt(i);
						}
						
						cppc.setAppPermission(methodName + "@" + lineNumber + "@" + reg + "@" + func + "@" + path);
					}
					// check Intent
					tss = tss[0].substring(1).split("[,]");
					if (func.equals("android.content.Context.registerReceiver")) {
						String i = tss[2].substring(1);
						Integer reg = Integer.parseInt(i);
						
						intpc.setAppPermission(methodName + "@" + lineNumber + "@" + reg + "@" + func + "@" + path);
					} else if (func.contains("sendBroadcast") ||
							   func.contains("sendOrderedBroadcast") ||
							   func.contains("sendStickyBroadcast") ||
							   func.contains("sendStickyOrderedBroadcast") ||
							   func.contains("startActivity") ||
							   func.contains("startActivityForResult") ||
							   func.contains("startService") ||
							   func.contains("bindService")) {
						String i = tss[1].substring(1);
						Integer reg = Integer.parseInt(i);
						
						intpc.setAppPermission(methodName + "@" + lineNumber + "@" + reg + "@" + func + "@" + path);
					}
				}
			}
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	public void getRegValueToCheck(String methodName, long lineNumber, int reg, String func, String path, int kind, PermissionChecker pc) {
		List<HashMap<String, String>> paramList = null;
		HashMap<Long, List<HashMap<String, String>>> ll = registerMapsForClasses.get(methodName);
		if (ll != null) {
			paramList = ll.get(lineNumber);
			boolean flag = false;
			for (HashMap<String, String> is : paramList) {
				String rv = is.get(String.valueOf(reg));
				if (rv != null) {
					if (rv.startsWith("VALUE(")) { // other class's field
						String field = rv.substring(6, rv.length() - 1);
						String oclassName = field.substring(0, field.lastIndexOf("."));
						HashMap<Long, List<HashMap<String, String>>> ohm = registerMapsForClasses.get(oclassName + ".<clinit>");
						if (ohm != null && !ohm.equals(ll)) { // avoid endless loop
							Iterator it = ohm.entrySet().iterator();
							while (it.hasNext()) {
								Map.Entry e1 = (Entry) it.next();
								List<HashMap<String, String>> regMapList = (List) e1.getValue();
								for (HashMap<String, String> rm : regMapList) {
									String fv = rm.get(field);
									if (fv != null && !fv.equals("")) {
										pc.analyse(permission_app_required, func, fv);
										flag = true;
									}
								}
							}
						}
					} else {
						pc.analyse(permission_app_required, func, rv);
						flag = true;
					}
				}
			}
			if (flag == false) {
				System.out.println(path);
				System.out.println(methodName);
				System.out.println(lineNumber);
				System.out.println("reg" + reg);
			}
		}
	}
	
	public static void printRegisterMapsForClasses() {
		if (registerMapsForClasses != null && registerMapsForClasses.size() > 0) {
			Iterator i = registerMapsForClasses.entrySet().iterator();
			while (i.hasNext()) {
				Map.Entry entry = (Map.Entry) i.next();
				if (!entry.getKey().toString().contains(
						"cn.edu.sjtu.test.Test.Simple_1"))
					continue;
				System.out.println("@ " + entry.getKey() + " @");
				HashMap<Long, List<HashMap<Integer, String>>> lis = (HashMap<Long, List<HashMap<Integer, String>>>) entry
						.getValue();
				Iterator j = lis.entrySet().iterator();
				while (j.hasNext()) {
					Map.Entry e = (Map.Entry) j.next();
					System.out.println("\t" + e.getKey());
					List<HashMap<Integer, String>> is = (List<HashMap<Integer, String>>) e
							.getValue();
					for (HashMap<Integer, String> it : is) {
						Iterator k = it.entrySet().iterator();
						while (k.hasNext()) {
							Map.Entry ee = (Map.Entry) k.next();
							System.out.println("\t\t" + ee.getKey() + ": "
									+ ee.getValue());
						}
						System.out.println("------------------------------");
					}
				}
			}
		} else {
			System.out.println("registerMapsForClasses is null or empty");
		}
	}

	public static void main(String[] args) {
		if (args.length < 4) {
			System.out.println("Missing arg...");
		} else {
			DedexerPlus dp = new DedexerPlus(args[0], args[1], args[2]);
			dp.checkAppPermission(args[3]);
		}
		printRegisterMapsForClasses();
	}
}
